# encoding :utf-8

import io  # 文件数据流
import datetime
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow import keras
# 导入常见网络层, sequential容器, 优化器, 损失函数
from tensorflow.keras import layers, Sequential, optimizers, losses, metrics
import os # 运维模块， 调用系统命令
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # 只显示warring和error


def preprocess(x, y):
    x = tf.cast(x, dtype=tf.float32) / 255.   #变成float，然后缩放
    y = tf.cast(y, dtype=tf.int32)
    return x, y

#将内存中图片转变成tf图片
def plot_to_image(figure):
    buf = io.BytesIO()  # 在内存中存储画
    plt.savefig(buf, format='png')
    plt.close(figure)
    buf.seek(0) # 存储修改为0
    # 传化为TF 图
    image = tf.image.decode_png(buf.getvalue(), channels=4) # (1000, 1000, 4)
    image = tf.expand_dims(image, 0)  # (1, 1000, 1000, 4)
    return image


def image_grid(images):
    # 返回一个5x5的mnist图像
    figure  = plt.figure(figsize=(10, 10))
    for i in range(25):
        plt.subplot(5, 5, i+1, title='name')
        plt.xticks([])
        plt.yticks([])
        plt.grid(False)
        plt.imshow(images[i], cmap=plt.cm.binary)
    return figure

#加载数据
batchsz = 128
(x, y), (x_val, y_val) = tf.keras.datasets.mnist.load_data()
print('datasets:', x.shape, y.shape, x.min(), x.max())

# 训练数据
# 将数据存储到数据管道中，提速
db = tf.data.Dataset.from_tensor_slices((x,y))
# 特征归一化，标签转成整数形式；shuffle洗牌；batch批次大小；repeat重复次数，即epoch=10
db = db.map(preprocess).shuffle(60000).batch(batchsz).repeat(10)  #10 epoch
# 验证数据
ds_val = tf.data.Dataset.from_tensor_slices((x_val, y_val))
ds_val = ds_val.map(preprocess).batch(batchsz, drop_remainder=True) # drop_remainder 清除不足一批次的数据

#定义模型
network = Sequential([
    layers.Dense(256, activation='relu'),
    layers.Dense(128, activation='relu'),
    layers.Dense(64, activation='relu'),
    layers.Dense(32, activation='relu'),
    layers.Dense(10)
])

# build  单独指定输入数据的维度
network.build(input_shape=(None, 28*28))  #28*28=784

network.summary()  #打印模型结构

#定义优化器
optimizer=optimizers.Adam(lr=0.01)
#指定tensorboard路径
current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S") #时间戳变成字符串
log_dir = 'logs1/' + current_time  #注意tf中的路径的写法 'logs1/' 或者 'logs1\\'
summary_writer = tf.summary.create_file_writer(log_dir)  # 创建监控类，监控数据写入到log_dir目录

sample_img = next(iter(db))[0] # 第一个批次的数据
sample_img = sample_img[0]  # 第一张图
sample_img = tf.reshape(sample_img, [1, 28, 28, 1])  #变成tf图片维度(样本个数1, 高度28,宽度28,通道1)
with summary_writer.as_default():  # 写入环境
    tf.summary.image("Training sample:", sample_img, step=0) #写入image

#梯度下降
for step, (x, y) in enumerate(db):    # 遍历切分好的数据step:0->599
    with tf.GradientTape() as tape:
        x = tf.reshape(x, (-1, 28*28))
        out = network(x)  # 正向传播--〉h
        y = tf.one_hot(y, depth=10) #独热编码
        loss = tf.reduce_mean(tf.losses.categorical_crossentropy(y, out, from_logits=True))  #需要独热

    grads = tape.gradient(loss, network.trainable_variables) #求梯度
    optimizer.apply_gradients(zip(grads, network.trainable_variables))  #更新参数

    if step % 100 == 0:
        print(step, 'loss:', float(loss))  # 读统计数据
        with summary_writer.as_default():
            tf.summary.scalar('train-loss', float(loss), step=step)  # 将loss写入到train-loss中

    if step % 500 == 0:
        total, total_correct = 0., 0
        # 计算准确率
        for _, (m, n) in enumerate(ds_val): # x, y
            m = tf.reshape(m, (-1, 28*28))
            out = network(m)
            pred = tf.argmax(out, axis=1)  #预测类别
            pred = tf.cast(pred, dtype=tf.int32)
            correct = tf.equal(pred, n)  #预测和真实对比
            total_correct += tf.reduce_sum(tf.cast(correct, dtype=tf.int32)).numpy()  #准确个数
            total += m.shape[0]  #总样本数

        print(step, 'Evaluate Acc:', total_correct / total)  #验证样本的准确率

        val_images = m[:25]  #取前25个图片
        val_images = tf.reshape(val_images, [-1, 28, 28, 1])  #转变成tf图片格式
        with summary_writer.as_default():
            tf.summary.scalar('test-acc', float(total_correct / total), step=step)  # 写入测试准确率
            tf.summary.image("val-onebyone-images:", val_images, max_outputs=25, step=step)  # 可视化测试用图片，25张
            val_images = tf.reshape(val_images, [-1, 28, 28])
            figure = image_grid(val_images) # image_grid 自定义函数：变成5*5网格图
            tf.summary.image('val-images:', plot_to_image(figure), step=step)  #转变成tf图片



'''
#autograph静态图中元素的tensorboard可视化也是相同的处理
#Example usage with tf.function graph execution:
#创建一个文件写入对象
writer = tf.summary.create_file_writer(logdir)

@tf.function
def my_func(step):
  # other model code would go here
  with writer.as_default():
    tf.summary.scalar("my_metric", 0.5, step=step)

for step in range(100):
  my_func(step)
  writer.flush()
  
'''